(* opt2.ml 
This optimisation remove unnecessary assignment and unused variable , for example :
...
v = 4
... if v isn't used here
v = 0

we can remove v = 4
*)
open List

open Variable
open Type

let in_if = ref false;;

(* list whish save variable used *)
let used_list : (string * (bool ref)) list ref = ref [] ;;

(* return value or NotFound *)
let rec get_value liste name = 
  match liste with
    [] -> raise NotFound
  | (name2, v)::tail ->  if ( compare name name2 == 0 ) then !v
              else get_value tail name
;;

let rec set_value liste name value =
  match liste with
    [] -> (used_list := (name, ref value)::!used_list)
  | (name2, v)::tail -> if ( compare name name2 == 0 ) then v := value
              else set_value tail name value
;;

let rec has_side_effect expr =
  match expr with
    Function(_,_) -> true
  | Binary(e1, op, e2) -> has_side_effect e1 && has_side_effect e2
  | Unary(op, e) -> has_side_effect e
  | _ -> false
;;
  


let rec remove_prev_assign name ast = 
  match ast with
    [] -> raise NotFound
  | Assignment(Var(name2),_)::tail
      when (compare name name2 == 0 ) -> tail
  | While(cond, body)::tail ->
      (try let new_body = remove_prev_assign name (List.rev body) in
         While(cond, rev new_body)::tail
              with NotFound ->
                While(cond, body)::remove_prev_assign name tail)
  | For(n, b, e, body)::tail ->
      (try let new_body = remove_prev_assign name (rev body) in
         For(n, b, e, rev new_body)::tail
              with NotFound ->
                For(n, b, e, body)::remove_prev_assign name tail)
  | Struct_If(if_list)::tail ->
      let new_if_list = remove_prev_assign_if name (rev if_list) in
        Struct_If(rev new_if_list)::remove_prev_assign name tail
  | head::tail -> head::remove_prev_assign name tail

and remove_prev_assign_if name expr =
  match expr with
    [] -> []
  | head::tail -> begin match head with 
            If(cond, body) ->
              (try let new_body = remove_prev_assign name (rev body) in
                  If(cond, rev new_body)::remove_prev_assign_if name tail
              with NotFound ->
                  If(cond, rev body)::remove_prev_assign_if name tail)
          | Else(body) ->
              (try let new_body = remove_prev_assign name (rev body) in
                  Else(rev new_body)::remove_prev_assign_if name tail
              with NotFound ->
                  Else(rev body)::remove_prev_assign_if name tail)
          end
;;

let rec remove_var ast =
  match ast with
    [] -> []
  | head::tail ->
    match head with
    Assignment(Var(name), e) ->
      let v = get_value !used_list name in
        if ( not v )
          then begin
          print_string ("  warning : unused variable `"^name^"`\n");
          set_value !used_list name true;
          remove_var tail
        end else head::(remove_var tail)
    | While(cond,body) -> While(cond, remove_var body)::(remove_var tail)
    | For(n,b,e,body) -> For(n,b,e, remove_var body)::(remove_var tail)
    | Struct_If(if_list) -> Struct_If(map remove_var_if if_list)
                              ::(remove_var tail)
    | _ -> head::(remove_var tail)

and remove_var_if expr =
  match expr with
    If(cond, body) -> If(cond, remove_var body)
  | Else(body) -> Else( remove_var body)

;;

let rec opt2_int next prev =
  match next with
  [] -> prev
  | head::tail ->
    begin match head with
      Assignment(Var(name), e) -> 
      print_string ("assign_var : "^name^" ");
      (* if e can possibly print something, e.g. here by calling a function *)
      if ( has_side_effect e)
        then begin
          print_string "side_effect !\n";
          set_value !used_list name true;
          let _ = opt2_int [e] [] in
          opt2_int tail (head::prev)
      (* don't remove ass out of if if we are in if statement *)
      end else if (!in_if) 
        then
        let _ = opt2_int [e] [] in
        opt2_int tail (head::prev)
      else begin
        let _ = opt2_int [e] [] in
        (try let v = get_value (!used_list) name in
          if ( v )
            then begin
            print_string "don't removed, used !\n";
            set_value !used_list name false;
            opt2_int tail (head::prev)
          end else begin
            print_string ("  remove assignment `"^name^"`\n");
            opt2_int tail (head::((remove_prev_assign name (prev))))
          end
        with NotFound ->
          print_string "first ass, don't remove \n";
          used_list := (name, ref false)::!used_list;
          opt2_int tail (head::prev))
      end
    | Var(name) -> print_string ("use_var :"^name) ;
                   set_value !used_list name true; opt2_int tail (head::prev)
    | Binary(e1, op, e2) -> 
      let _ = opt2_int [e1] [] and
      _ = opt2_int [e2] [] in
      opt2_int tail (head::prev)
    | Unary(op, e) ->
      let new_e = opt2_int [e] [] in
      opt2_int tail (Unary(op, hd new_e)::prev)
    | Function(name, args) -> let _ = opt2_int args [] in
                                opt2_int tail (head::prev)
    | For(var, beg, end_, body) ->
      let _ = opt2_int [var] [] in
      let new_body = opt2_int body [] in
      let _ = opt2_int [var] [] in
      opt2_int tail (For(var, beg, end_, rev new_body)::prev)
    | While(cond, body) -> 
      let _ = opt2_int [cond] [] in
      let new_body = opt2_int body [] in
      let _ = opt2_int [cond] [] in
      opt2_int tail ((While(cond, rev new_body))::prev)
    | Struct_If(if_list) ->
      in_if := true;
      let new_if_list = Struct_If(map opt2_int_if if_list) in
      in_if := false;
      opt2_int tail (new_if_list::prev)
    | Print(plist) -> let _ = opt2_int plist [] in
      opt2_int tail (head::prev)
    | Int(_) -> opt2_int tail (head::prev)
    | Float(_) -> opt2_int tail (head::prev)
    | Bool(_) -> opt2_int tail (head::prev)
    | String(_) -> opt2_int tail (head::prev)
    | _ -> opt2_int tail (head::prev)
    end

and opt2_int_if expr = 
  match expr with
      If(cond, body) ->let _ = opt2_int [cond] [] in If(cond, opt2_int body [])
    | Else(body) -> Else(opt2_int body [])
;;

let opt2 ast = 
  let new_ast = opt2_int ast [] in rev (remove_var new_ast)
;;